package cn.edu.bupt.fluxdemo.service;

import org.springframework.stereotype.Component;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.util.Queue;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;

@Component
public class UserMessageService {
    // 存储用户消息队列：Key=用户ID，Value=消息队列
    private final ConcurrentHashMap<String, Queue<String>> userMessages = new ConcurrentHashMap<>();

    // 添加消息到指定用户的队列
    public void addMessage(String userId, String message) {
        userMessages.computeIfAbsent(userId, k -> new ConcurrentLinkedQueue<>()).add(message);
    }

    // 监听用户消息（非阻塞轮询）
    public Mono<String> waitForMessage(String userId, Duration timeout) {
        return Mono.create(sink -> {
            Queue<String> queue = userMessages.computeIfAbsent(userId, k -> new ConcurrentLinkedQueue<>());
            // 立即检查是否有消息
            String message = queue.poll();
            if (message != null) {
                sink.success(message); // 有消息，立即返回
                return;
            }
            // 无消息，启动后台线程监听
            ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
            AtomicBoolean responded = new AtomicBoolean(false);
            // 每隔1秒检查一次消息
            executor.scheduleAtFixedRate(() -> {
                String msg = queue.poll();
                if (msg != null && !responded.get()) {
                    responded.set(true);
                    sink.success(msg); // 监听到消息，返回
                    executor.shutdown();
                }
            }, 0, 1, TimeUnit.SECONDS);
            // 设置超时
            executor.schedule(() -> {
                if (!responded.get()) {
                    responded.set(true);
                    sink.success(""); // 超时返回空
                    executor.shutdown();
                }
            }, timeout.toMillis(), TimeUnit.MILLISECONDS);
        });
    }
}